import numpy as np
import torch
from sklearn.cluster import KMeans
import os
import yaml

def load_annotations(dataset_path):
    """加载数据集标注，提取边界框信息"""
    bboxes = []
    labels_path = os.path.join(dataset_path, 'labels')
    
    for label_file in os.listdir(labels_path):
        if label_file.endswith('.txt'):
            with open(os.path.join(labels_path, label_file), 'r') as f:
                lines = f.readlines()
                
            for line in lines:
                parts = line.strip().split()
                if len(parts) >= 5:
                    # YOLO格式: class x_center y_center width height
                    w = float(parts[3])
                    h = float(parts[4])
                    bboxes.append([w, h])
                     
    return np.array(bboxes)


def calculate_adaptive_anchors(dataset_path, num_anchors=9):
    """使用K-means聚类计算自适应锚框"""
    # 加载边界框数据
    bboxes = load_annotations(dataset_path)
    if len(bboxes) == 0:
        raise ValueError("未找到边界框数据，请检查数据集路径和格式")
    
    # 使用K-means聚类计算锚框
    kmeans = KMeans(n_clusters=num_anchors, random_state=0)
    kmeans.fit(bboxes)
    anchors = kmeans.cluster_centers_
    
    # 按面积排序锚框
    areas = anchors[:, 0] * anchors[:, 1]
    sorted_indices = np.argsort(areas)
    anchors = anchors[sorted_indices]
    
    # 转换为YOLOv5格式 (wh) 并四舍五入
    anchors = np.round(anchors * 640).astype(int).tolist()
    
    print(f"计算得到的自适应锚框: {anchors}")
    return anchors


def update_anchors_in_config(config_path, new_anchors):
    """更新配置文件中的锚框参数"""
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    
    config['anchors'] = new_anchors
    
    with open(config_path, 'w') as f:
        yaml.safe_dump(config, f, sort_keys=False)
    
    print(f"已更新配置文件 {config_path} 中的锚框参数")


def anchor_fitness(anchors, targets):
    """计算锚框适应度，用于评估锚框质量"""
    # 计算锚框与目标的交并比
    anchors = torch.tensor(anchors, dtype=torch.float32)
    targets = torch.tensor(targets, dtype=torch.float32)
    
    # 计算交并比
    intersection = torch.min(anchors[:, None, 0], targets[:, 0]) * torch.min(anchors[:, None, 1], targets[:, 1])
    union = anchors[:, None, 0] * anchors[:, None, 1] + targets[:, 0] * targets[:, 1] - intersection
    iou = intersection / union  # [num_anchors, num_targets]
    
    # 计算适应度 (考虑召回率和平均IoU)
    best_iou, _ = iou.max(0)
    fitness = (best_iou * (1 - best_iou)).mean()
    
    return fitness.item()

if __name__ == '__main__':
    # 示例用法
    dataset_path = '../data/train'
    config_path = '../data/config.yaml'
    anchors = calculate_adaptive_anchors(dataset_path)
    update_anchors_in_config(config_path, anchors)